import csv
import os
import glob

# 输入文件夹和输出文件夹的路径
input_folder = '/root/autodl-tmp/TED/TEDLIUM_release1/train/stm'
output_folder = '/root/autodl-tmp/TEDspliter/train/stm'

# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)

# 获取所有 STM 文件
stm_files = glob.glob(os.path.join(input_folder, '*.stm'))

# 对于每一个 STM 文件
for stm_file in stm_files:
    # 获取输入文件的基本名称（无路径）
    base_name = os.path.basename(stm_file)
    # 创建输出文件的路径
    output_file = os.path.join(output_folder, os.path.splitext(base_name)[0] + '.csv')
    
    # 打开输入文件和输出文件
    with open(stm_file, 'r') as infile, open(output_file, 'w', newline='') as outfile:
        # 创建 csv 写入器
        writer = csv.writer(outfile)
        # 写入标题行
        writer.writerow(['xmin', 'xmax', 'caption'])
        
        # 对于输入文件中的每一行
        for line in infile:
            # 分割行
            parts = line.split()
            
            # 提取所需的部分
            xmin = parts[3]
            xmax = parts[4]
            caption = ' '.join(parts[6:-1])
            
            # 写入一行到输出文件
            writer.writerow([xmin, xmax, caption])
